0%

SpringAI — 向量数据库能力

概述

向量数据库是一种特殊类型的数据库,在 AI 应用中扮演着至关重要的角色。在向量数据库中,查询方式与传统关系型数据库不同,向量数据库执行的是相似度搜索,而非精确匹配。当给定一个向量作为查询时,向量数据库会返回与查询向量”相似”的向量。

向量数据库用于将您的数据与 AI 模型集成,其使用步骤如下:

  1. 首先将您的数据加载到向量数据库中。
  2. 当用户查询需要发送给 AI 模型时,先检索一组相似的文档。
  3. 这些文档随后作为用户问题的上下文,与用户查询一起发送给 AI 模型。

这种技术被称为检索增强生成(Retrieval Augmented Generation,RAG)

本文将介绍 Spring AI 中使用多种向量数据库实现的接口,以及一些高层级的使用示例。最后一节旨在揭开向量数据库中相似度搜索底层方法的神秘面纱。


VectorStore 接口

Spring AI 通过 VectorStore 接口及其只读对应接口 VectorStoreRetriever,提供了与向量数据库交互的抽象 API。

VectorStoreRetriever — 只读接口

Spring AI 提供名为 VectorStoreRetriever 的只读接口,仅暴露文档检索功能:

1
2
3
4
public interface VectorStoreRetriever {
List<Document> similaritySearch(String query);
List<Document> similaritySearch(SearchRequest request);
}

这个函数式接口专为只需要从向量存储中检索文档而无需执行任何变更操作的场景而设计,它通过仅暴露文档检索所需的功能,遵循了最小权限原则

VectorStore — 读写接口

VectorStore 接口扩展了 VectorStoreRetriever,并增加了变更操作的能力:

1
2
3
4
5
public interface VectorStore extends VectorStoreRetriever {
void add(List<Document> documents);
void delete(List<String> idList);
void delete(Filter.Expression filterExpression);
}

VectorStore 接口结合了读取和写入操作,允许您在向量数据库中添加、删除和搜索文档。

插入数据

要向向量数据库插入数据,需要将其封装在 Document 对象中。Document 类封装了来自数据源(如 PDF 或 Word 文档)的内容,包含以字符串形式表示的文本,同时还包含以键值对形式存储的元数据,包括文件名等详细信息。

插入向量数据库后,文本内容会使用嵌入模型转换为数值数组(即 float[]),称为向量嵌入(vector embeddings)。嵌入模型(如 Word2Vec、GLoVE、BERT 或 OpenAI 的 text-embedding-ada-002)用于将单词、句子或段落转换为这些向量嵌入。

向量数据库的作用是存储这些嵌入并执行相似度搜索,它本身不生成嵌入,要创建向量嵌入,应使用 EmbeddingModel

相似度搜索参数

接口中的 similaritySearch 方法允许检索与给定查询字符串相似的文档。这些方法可以通过以下参数进行微调:

参数 说明
k 一个整数,指定要返回的最大相似文档数量。这通常被称为”top K”搜索或”K 近邻”(KNN)搜索。
threshold 一个介于 0 到 1 之间的双精度值,值越接近 1 表示相似度越高。例如,如果设置阈值为 0.75,则仅返回相似度高于此值的文档。
Filter.Expression 一个用于传递流畅 DSL(领域特定语言)表达式的类,功能类似于 SQL 中的 WHERE 子句,但仅适用于 Document 的元数据键值对。
filterExpression 基于 ANTLR4 的外部 DSL,接受字符串形式的过滤表达式。例如,对于 countryyearisActive 等元数据键,可以使用如下表达式:country == 'UK' && year >= 2020 && isActive == true

模式初始化

某些向量存储在使用前需要初始化其后端模式,默认情况下不会自动为您初始化。您必须选择启用——通过向适当的构造函数参数传递一个布尔值,或者如果使用的是 Spring Boot,则在 application.propertiesapplication.yml 中将相应的 initialize-schema 属性设置为 true。请查阅您所使用的向量存储的文档以获取具体的属性名称。


批处理策略(Batching Strategy)

为什么需要批处理?

在使用向量存储时,通常需要嵌入大量文档。虽然一次性调用嵌入所有文档看似简单,但这种方法可能会导致问题。嵌入模型以 token 为单位处理文本,并且有最大 token 限制(通常称为上下文窗口大小)。这个限制约束了单次嵌入请求中可以处理的文本量,尝试在一次调用中嵌入过多 token 可能会导致错误或截断嵌入。

TokenCountBatchingStrategy

为了解决上述问题,Spring AI 实现了批处理策略,这种方法将大批量文档拆分为符合嵌入模型最大上下文窗口的小批次。批处理不仅解决了 token 限制问题,还可以提高性能并更有效地利用 API 速率限制。

Spring AI 通过 BatchingStrategy 接口提供此功能,允许根据 token 数量以子批次方式处理文档。

核心 BatchingStrategy 接口定义如下:

定义了一个单一方法 batch,它接受文档列表并返回文档批次列表

1
2
3
public interface BatchingStrategy {
List<List<Document>> batch(List<Document> documents);
}

Spring AI 提供了一个名为 TokenCountBatchingStrategy 的默认实现,该策略根据文档的 token 数量进行批处理,确保每个批次不超过计算出的最大输入 token 数量。

TokenCountBatchingStrategy 的关键特性:

  • 包含预留百分比(默认 10%),为潜在的开销提供缓冲。
  • 实际最大输入 token 数量的计算公式为:actualMaxInputTokenCount = originalMaxInputTokenCount * (1 - RESERVE_PERCENTAGE)
  • 该策略会估算每个文档的 token 数量,在不超过最大输入 token 数量的前提下将它们分组为批次;如果单个文档超过此限制,则抛出异常。

自定义 TokenCountBatchingStrategy

您也可以自定义 TokenCountBatchingStrategy 以更好地满足特定需求,可以在 Spring Boot 的 @Configuration 类中创建具有自定义参数的新实例。

以下是创建自定义 TokenCountBatchingStrategy Bean 的示例:

1
2
3
4
5
6
7
8
9
10
11
12
@Configuration
public class BatchingConfig {

@Bean
public BatchingStrategy customBatchingStrategy() {
return new TokenCountBatchingStrategy(
EncodingType.CL100K_BASE, // 用于 token 化的编码类型
8000, // 最大输入 token 数量
0.1 // 预留百分比
);
}
}

在此配置中:

  • EncodingType.CL100K_BASE:指定用于 token 化的编码类型。JTokkitTokenCountEstimator 使用此编码类型来准确估算 token 数量。
  • 8000:设置最大输入 token 数量。此值应小于或等于嵌入模型的最大上下文窗口大小。
  • 0.1:设置预留百分比。从最大输入 token 数量中预留的 token 百分比,为处理过程中潜在的 token 数量增加创建缓冲。

默认情况下,此构造函数使用 Document.DEFAULT_CONTENT_FORMATTER 进行内容格式化,并使用 MetadataMode.NONE 进行元数据处理。如果需要自定义这些参数,可以使用带有额外参数的完整构造函数。定义后,此自定义 TokenCountBatchingStrategy Bean 将自动被应用程序中的 EmbeddingModel 实现使用,替换默认策略。

TokenCountBatchingStrategy 内部使用 TokenCountEstimator(具体为 JTokkitTokenCountEstimator)来计算 token 数量,以实现高效批处理。这确保了基于指定编码类型的准确 token 估算。此外,TokenCountBatchingStrategy 允许您传入自己实现的 TokenCountEstimator 接口,从而提供灵活性。此功能使您能够使用针对特定需求定制的自定义 token 计数策略。

自动截断(Auto-Truncation)场景

某些嵌入模型(如 Vertex AI 文本嵌入)支持 auto_truncate 功能。启用后,模型会静默截断超过最大大小的文本输入并继续处理;禁用时,会对过大的输入抛出显式错误。

当将批处理策略与自动截断功能结合使用时,必须将批处理策略配置为使用远高于模型实际最大值的输入 token 数量。这可以防止批处理策略对大文档引发异常,从而允许嵌入模型在内部处理截断。

启用自动截断时,请将批处理策略的最大输入 token 数量设置为远高于模型实际限制的值。

以下是一个在 PgVectorStore 中使用 Vertex AI 自动截断和自定义 BatchingStrategy 的配置示例:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
@Configuration
public class VectorStoreConfig {

@Bean
public EmbeddingModel embeddingModel() {
// 启用自动截断的 Vertex AI 嵌入模型
return new VertexAiEmbeddingModel(
"text-embedding-004",
true // auto-truncate enabled
);
}

@Bean
public BatchingStrategy customBatchingStrategy() {
// 使用远高于实际模型限制 (20,000) 的 token 限制 (132,900)
return new TokenCountBatchingStrategy(
EncodingType.CL100K_BASE,
132900,
0.1
);
}

@Bean
public VectorStore vectorStore(EmbeddingModel embeddingModel, BatchingStrategy batchingStrategy) {
return PgVectorStore.builder(embeddingModel)
.batchingStrategy(batchingStrategy)
.build();
}
}

在此配置中:

  • 嵌入模型启用了自动截断,使其能够优雅地处理超大输入。
  • 批处理策略使用了人为调高的 token 限制(132,900),远高于实际模型限制(20,000)。
  • 向量存储使用配置好的嵌入模型和自定义 BatchingStrategy Bean。

这种方法之所以有效,是因为:

  1. TokenCountBatchingStrategy 检查是否有单个文档超过配置的最大值,如果有则抛出 IllegalArgumentException
  2. 通过在批处理策略中设置非常高的限制,我们确保此检查永远不会失败。
  3. 超过模型限制的文档或批次会被嵌入模型的自动截断功能静默截断并处理。

使用自动截断时的注意事项:

  • 将批处理策略的最大输入 token 数量设置为模型实际限制的至少 5-10 倍,以避免批处理策略过早抛出异常。
  • 监控日志中来自嵌入模型的截断警告(注意:并非所有模型都会记录截断事件)。
  • 考虑静默截断对嵌入质量的影响。
  • 使用样本文档进行测试,确保截断后的嵌入仍然满足您的需求。
  • 为未来的维护者记录此配置,因为这是非标准做法。

Spring Boot 自动配置

如果您使用的是 Spring Boot 自动配置,则必须提供一个自定义的 BatchingStrategy Bean 来覆盖 Spring AI 自带的默认策略:

1
2
3
4
5
6
7
8
9
10
11
12
@Configuration
public class BatchingConfig {

@Bean
public BatchingStrategy batchingStrategy() {
return new TokenCountBatchingStrategy(
EncodingType.CL100K_BASE,
8000,
0.1
);
}
}

此 Bean 存在于应用程序上下文中后,将自动替换所有向量存储使用的默认批处理策略。


VectorStore 实现

以下是 VectorStore 接口的可用实现:

实现 说明
SimpleVectorStore 一个简单的向量存储实现,仅适用于测试目的
更多实现 点击查看

使用示例

选择嵌入模型

要为向量数据库计算嵌入,您需要选择与所使用的高层 AI 模型相匹配的嵌入模型。

例如,使用 OpenAI 的 ChatGPT 时,我们使用 OpenAiEmbeddingModel 和名为 text-embedding-ada-002 的模型。

OpenAI 的 Spring Boot Starter 自动配置使得 EmbeddingModel 的实现可在 Spring 应用程序上下文中进行依赖注入。

加载数据

将数据加载到向量存储的一般用法类似于批处理作业:首先将数据加载到 Spring AI 的 Document 类中,然后调用 VectorStore 接口上的 add 方法。

给定一个表示 JSON 文件的源文件 String 引用,我们使用 Spring AI 的 JsonReader 加载 JSON 中的特定字段,将其拆分为小块,然后将这些小块传递给向量存储实现。VectorStore 实现会计算嵌入并将 JSON 和嵌入一起存储在向量数据库中:

1
2
3
4
5
6
7
8
@Autowired
private VectorStore vectorStore;

public void loadData() {
JsonReader reader = new JsonReader("data.json", "content");
List<Document> documents = reader.get();
vectorStore.add(documents);
}

随后,当用户的提问传递给 AI 模型时,会执行相似度搜索以检索相似的文档,然后这些文档被”填充”到提示中作为用户问题的上下文。

检索数据

对于只读操作,您可以使用 VectorStore 接口或更专注的 VectorStoreRetriever 接口:

1
2
3
4
5
6
7
8
// 使用 VectorStore 进行搜索
List<Document> results = vectorStore.similaritySearch("What is AI?");

// 使用 VectorStoreRetriever 进行搜索(只读)
@Autowired
private VectorStoreRetriever vectorStoreRetriever;

List<Document> results = vectorStoreRetriever.similaritySearch("What is AI?");

可以向 similaritySearch 方法传递额外选项,以定义要检索的文档数量和相似度搜索的阈值:

1
2
3
4
5
6
7
SearchRequest request = SearchRequest.builder()
.query("What is AI?")
.topK(5)
.similarityThreshold(0.75)
.build();

List<Document> results = vectorStore.similaritySearch(request);

使用单独的接口可以让您清晰地定义哪些组件需要写入权限,哪些只需要读取权限:

1
2
3
4
5
6
7
8
9
10
11
12
@Service
public class DocumentRetrievalService {
private final VectorStoreRetriever retriever;

public DocumentRetrievalService(VectorStoreRetriever retriever) {
this.retriever = retriever;
}

public List<Document> findRelevantDocs(String query) {
return retriever.similaritySearch(query);
}
}

这种关注点分离通过仅向真正需要它们的组件开放变更操作,有助于创建更可维护和更安全的应用程序。

VectorStoreRetriever 的优势

VectorStoreRetriever 接口提供了向量存储的只读视图,仅暴露相似度搜索功能。它遵循最小权限原则,在只需要检索文档而无需修改底层数据的 RAG(检索增强生成)应用中特别有用。

优势 说明
关注点分离 清晰地将读取操作与写入操作分离。
接口隔离 只需要检索功能的客户端不会暴露变更方法。
函数式接口 可以使用 lambda 表达式或方法引用实现简单用例。
减少依赖 只需要执行搜索的组件不需要依赖完整的 VectorStore 接口。

RAG 应用中的使用

VectorStoreRetriever 接口在 RAG 应用中特别有用,您需要检索相关文档以为 AI 模型提供上下文:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
@Service
public class RagService {
private final VectorStoreRetriever retriever;
private final ChatModel chatModel;

public RagService(VectorStoreRetriever retriever, ChatModel chatModel) {
this.retriever = retriever;
this.chatModel = chatModel;
}

public String answer(String question) {
// 检索相关文档
List<Document> relevantDocs = retriever.similaritySearch(question);

// 构建包含上下文的提示
String context = relevantDocs.stream()
.map(Document::getText)
.collect(Collectors.joining("\n"));

String prompt = "基于以下上下文回答问题:\n" + context + "\n\n问题:" + question;

// 生成回答
return chatModel.call(prompt);
}
}

这种模式允许在 RAG 应用中清晰地分离检索组件和生成组件。


过滤表达式

这里来介绍可用于查询结果的各类过滤器。您可以将类似 SQL 的过滤表达式作为 String 传递给 similaritySearch 的重载方法之一。

字符串过滤表达式示例

考虑以下示例:

1
2
3
"country == 'BG'"
"genre == 'drama' && year >= 2020"
"genre in ['comedy', 'documentary', 'drama']"

FilterExpressionBuilder — 流畅 API

您可以使用 FilterExpressionBuilder 创建 Filter.Expression 的实例,它提供了一个流畅的 API。

简单示例如下:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
FilterExpressionBuilder builder = new FilterExpressionBuilder();

// country == 'BG'
Filter.Expression expr1 = builder.eq("country", "BG").build();

// genre == 'drama' && year >= 2020
Filter.Expression expr2 = builder
.eq("genre", "drama")
.and(builder.gte("year", 2020))
.build();

// genre in ['comedy', 'documentary', 'drama']
Filter.Expression expr3 = builder
.in("genre", "comedy", "documentary", "drama")
.build();

可用操作符

您可以使用以下操作符构建复杂的表达式:

比较操作符:

操作符 方法 示例
== eq() builder.eq("country", "UK")
!= ne() builder.ne("status", "inactive")
> gt() builder.gt("year", 2020)
>= gte() builder.gte("year", 2020)
< lt() builder.lt("price", 100)
<= lte() builder.lte("price", 100)
in in() builder.in("genre", "comedy", "drama")

逻辑操作符:

操作符 方法 示例
&& and() builder.eq("a", 1).and(builder.eq("b", 2))
` `
! not() builder.not(builder.eq("status", "deleted"))

删除文档

Vector Store 接口提供了多种删除文档的方法,允许您通过特定的文档 ID 或使用过滤表达式来移除数据。

按 ID 列表删除

删除文档最简单的方法是提供文档 ID 列表:

1
2
List<String> idsToDelete = Arrays.asList("doc1", "doc2", "doc3");
vectorStore.delete(idsToDelete);

此方法会移除 ID 与提供列表中相匹配的所有文档,如果列表中的任何 ID 在存储中不存在,则会被忽略。

按过滤表达式删除

对于更复杂的删除条件,您可以使用过滤表达式:

1
2
3
4
5
6
7
// 使用 Filter.Expression
Filter.Expression filter = new FilterExpressionBuilder()
.eq("tenantId", "tenant-123")
.and(new FilterExpressionBuilder().lt("version", 2))
.build();

vectorStore.delete(filter);

此方法接受一个 Filter.Expression 对象,该对象定义了应删除哪些文档的条件。当您需要根据文档的元数据属性删除文档时,它特别有用。

按字符串过滤表达式删除

为方便起见,您也可以使用基于字符串的过滤表达式删除文档:

1
vectorStore.delete("tenantId == 'tenant-123' && version < 2");

此方法在内部将提供的字符串过滤器转换为 Filter.Expression 对象。当您以字符串格式拥有过滤条件时,这非常有用。

异常处理

所有删除方法在出现错误时都可能抛出异常:

1
2
3
4
5
6
try {
vectorStore.delete(idsToDelete);
} catch (Exception e) {
// 处理删除失败
log.error("Failed to delete documents", e);
}

最佳实践是将删除操作包装在 try-catch 块中:

1
2
3
4
5
6
7
8
9
public void safeDeleteDocuments(List<String> ids) {
try {
vectorStore.delete(ids);
log.info("Successfully deleted {} documents", ids.size());
} catch (Exception e) {
log.error("Failed to delete documents: {}", ids, e);
throw new DocumentDeletionException("Could not delete documents", e);
}
}

常见场景:管理文档版本

一个常见的场景是管理文档版本,您需要上传文档的新版本同时移除旧版本。以下是如何使用过滤表达式处理此问题:

1
2
3
4
5
6
7
8
9
10
// 上传新版本前,删除旧版本
Filter.Expression oldVersionFilter = new FilterExpressionBuilder()
.eq("documentId", "doc-456")
.and(new FilterExpressionBuilder().ne("version", newVersion))
.build();

vectorStore.delete(oldVersionFilter);

// 然后添加新版本
vectorStore.add(newDocument);

您也可以使用字符串过滤表达式完成同样的操作:

1
vectorStore.delete("documentId == 'doc-456' && version != " + newVersion);

删除最佳实践

  • 按 ID 列表删除在您确切知道要移除哪些文档时通常更快。
  • 基于过滤器的删除可能需要扫描索引以查找匹配的文档;不过,这取决于具体的向量存储实现。
  • 大型删除操作应进行批处理,以避免压垮系统。
  • 当基于文档属性而非先收集 ID 进行删除时,考虑使用过滤表达式。